package com.example.demo.baidu;

import com.baidu.aip.bodyanalysis.AipBodyAnalysis;
import com.baidu.aip.util.Base64Util;
import com.example.demo.baidu.utils.FileUtil;
import com.example.demo.baidu.utils.HttpUtil;
import org.json.JSONObject;

import javax.imageio.ImageIO;
import java.awt.*;
import java.awt.image.BufferedImage;
import java.io.*;
import java.net.URLEncoder;
import java.util.Base64;
import java.util.HashMap;

/**
 * @author void
 * @date 2021/6/3 16:44
 * @desc
 */
public class BodySeg {

    public static void api(){
        String url = "https://aip.baidubce.com/rest/2.0/image-classify/v1/body_seg";
        try {
            // 本地文件路径
            String filePath = "G://QQ图片20201209234847.jpg";
            byte[] imgData = FileUtil.readFileByBytes(filePath);
            String imgStr = Base64Util.encode(imgData);
            String imgParam = URLEncoder.encode(imgStr, "UTF-8");

            String param = "image=" + imgParam;

            // 注意这里仅为了简化编码每一次请求都去获取access_token，线上环境access_token有过期时间， 客户端可自行缓存，过期后重新获取。
            String accessToken = BaiduConstant.ACCESS_TOKEN;

            String result = HttpUtil.post(url, accessToken, param);
            JSONObject jsonObject = new JSONObject(result);
            byte[] foreground = Base64Util.decode(jsonObject.getString("foreground"));
            byte[] scoremap = Base64Util.decode(jsonObject.getString("scoremap"));

            FileOutputStream fout1 = new FileOutputStream("G://foreground.png");
            //将字节写入文件
            fout1.write(foreground);
            fout1.close();

            FileOutputStream fout2 = new FileOutputStream("G://scoremap.png");
            //将字节写入文件
            fout2.write(scoremap);
            fout2.close();
        } catch (Exception e) {
            e.printStackTrace();
        }
    }

    public static void main(String[] args) {
//        api();
        sdk();
    }

    public static void sdk(){
        // 初始化一个AipBodyAnalysis
        AipBodyAnalysis client = new AipBodyAnalysis(BaiduConstant.APP_ID, BaiduConstant.API_KEY, BaiduConstant.SECRET_KEY);

        HashMap<String, String> options = new HashMap<>();
        options.put("type", "foreground");

        // 可选：设置网络连接参数
        client.setConnectionTimeoutInMillis(2000);
        client.setSocketTimeoutInMillis(60000);

        // 调用接口
        String path = "G://QQ图片20201209234847.jpg";
        long d = System.currentTimeMillis();
        JSONObject res = client.bodySeg(path, options);
        System.out.println("耗时:"+(System.currentTimeMillis()-d));
        System.out.println(res.toString(2));

        FileOutputStream fout1;
        try {
            fout1 = new FileOutputStream("G://foreground.png");
            //将字节写入文件
            byte[] foreground = Base64Util.decode(res.getString("foreground"));
            fout1.write(foreground);
            fout1.close();
        } catch (FileNotFoundException e) {
            e.printStackTrace();
        } catch (IOException e) {
            e.printStackTrace();
        }
    }

    public static BufferedImage resize(BufferedImage img, int newW, int newH) {
        Image tmp = img.getScaledInstance(newW, newH, Image.SCALE_SMOOTH);
        BufferedImage dimg = new BufferedImage(newW, newH, BufferedImage.TYPE_INT_ARGB);

        Graphics2D g2d = dimg.createGraphics();
        g2d.drawImage(tmp, 0, 0, null);
        g2d.dispose();

        return dimg;
    }

    public static void convert(String labelmapBase64, int realWidth, int realHeight) {
        try {
            byte[] bytes = Base64.getDecoder().decode(labelmapBase64);
            InputStream is = new ByteArrayInputStream(bytes);
            BufferedImage image = ImageIO.read(is);
            BufferedImage newImage = resize(image, realWidth, realHeight);
            BufferedImage grayImage = new BufferedImage(realWidth, realHeight, BufferedImage.TYPE_BYTE_GRAY);
            for(int i= 0 ; i < realWidth ; i++){
                for(int j = 0 ; j < realHeight; j++){
                    int rgb = newImage.getRGB(i, j);
                    grayImage.setRGB(i, j, rgb * 255);  //将像素存入缓冲区
                }
            }
            File newFile = new File("gray.jpg");
            ImageIO.write(grayImage, "jpg", newFile);
        } catch (IOException e) {
            e.printStackTrace();
        }
    }
}
